package com.sie.algorithm.sm;

import java.math.BigInteger;
import java.security.SecureRandom;

import com.sie.algorithm.crypto.AsymmetricCipherKeyPair;
import com.sie.algorithm.crypto.digests.SM3Digest;
import com.sie.algorithm.crypto.generators.ECKeyPairGenerator;
import com.sie.algorithm.crypto.params.ECDomainParameters;
import com.sie.algorithm.crypto.params.ECKeyGenerationParameters;
import com.sie.algorithm.crypto.params.ECPrivateKeyParameters;
import com.sie.algorithm.crypto.params.ECPublicKeyParameters;
import com.sie.algorithm.math.ec.ECCurve;
import com.sie.algorithm.math.ec.ECFieldElement;
import com.sie.algorithm.math.ec.ECPoint;
import com.sie.zhq.tool.ByteTool;
import com.sie.zhq.tool.StringUtil;

/**
 * SM2加解密类
 *
 * @author zhq
 * @version 2016-04-01
 *
 */
public class SM2 {

    private int ct = 1;
    private ECPoint p2;
    private SM3Digest sm3keybase;
    private SM3Digest sm3c3;

    private byte[] key = new byte[32];
    private byte keyOff = 0;
    public static BigInteger ecc_p;
    public static BigInteger ecc_a;
    public static BigInteger ecc_b;
    public static BigInteger ecc_n;
    public static BigInteger ecc_gx;
    public static BigInteger ecc_gy;

    public static ECCurve ecc_curve;
    public static ECPoint ecc_point_g;

    public static ECDomainParameters ecc_bc_spec;
    public static ECKeyPairGenerator eccKeyPairGenerator;

    public ECFieldElement ecc_gx_fieldelement;
    public ECFieldElement ecc_gy_fieldelement;

    public static final String[] sm2Param = {
            "FFFFFFFEFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF00000000FFFFFFFFFFFFFFFF",// p,0
            "FFFFFFFEFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF00000000FFFFFFFFFFFFFFFC",// a,1
            "28E9FA9E9D9F5E344D5A9E4BCF6509A7F39789F515AB8F92DDBCBD414D940E93",// b,2
            "FFFFFFFEFFFFFFFFFFFFFFFFFFFFFFFF7203DF6B21C6052B53BBF40939D54123",// n,3
            "32C4AE2C1F1981195F9904466A39C9948FE30BBFF2660BE1715A4589334C74C7",// gx,4
            "BC3736A2F4F6779C59BDCEE36B692153D0A9877CC62A474002DF32E52139F0A0" // gy,5
    };

    public SM2() {

        ecc_p = new BigInteger(sm2Param[0], 16);
        ecc_a = new BigInteger(sm2Param[1], 16);
        ecc_b = new BigInteger(sm2Param[2], 16);
        ecc_n = new BigInteger(sm2Param[3], 16);
        ecc_gx = new BigInteger(sm2Param[4], 16);
        ecc_gy = new BigInteger(sm2Param[5], 16);

        ecc_curve = new ECCurve.Fp(ecc_p, ecc_a, ecc_b);
        ecc_point_g = ecc_curve.createPoint(ecc_gx, ecc_gy);
        // System.out.println("g:"+ecc_point_g.toString());
        ecc_bc_spec = new ECDomainParameters(ecc_curve, ecc_point_g, ecc_n);
        ECKeyGenerationParameters ecc_ecgenparam;
        ecc_ecgenparam = new ECKeyGenerationParameters(ecc_bc_spec,
                new SecureRandom());
        eccKeyPairGenerator = new ECKeyPairGenerator();
        eccKeyPairGenerator.init(ecc_ecgenparam);

    }

    private void reset() {
        this.sm3keybase = new SM3Digest();
        this.sm3c3 = new SM3Digest();
        byte p[] = ByteTool.byteConvert32Bytes(p2.normalize().getXCoord()
                .toBigInteger());
        this.sm3keybase.update(p, 0, p.length);
        this.sm3c3.update(p, 0, p.length);
        p = ByteTool.byteConvert32Bytes(p2.normalize().getYCoord()
                .toBigInteger());
        this.sm3keybase.update(p, 0, p.length);
        this.ct = 1;
        nextKey();
    }

    private void nextKey() {
        SM3Digest sm3keycur = new SM3Digest(sm3keybase);
        sm3keycur.update((byte) (ct >> 24 & 0x00ff));
        sm3keycur.update((byte) (ct >> 16 & 0x00ff));
        sm3keycur.update((byte) (ct >> 8 & 0x00ff));
        sm3keycur.update((byte) (ct & 0x00ff));
        sm3keycur.doFinal(key, 0);
        keyOff = 0;
        ct++;
    }

    public SM2keyPair generateKeyPair() {
        BigInteger k = null;
        ECPoint c1 = null;
        AsymmetricCipherKeyPair key = eccKeyPairGenerator.generateKeyPair();
        ECPrivateKeyParameters ecpriv = (ECPrivateKeyParameters) key
                .getPrivate();
        ECPublicKeyParameters ecpub = (ECPublicKeyParameters) key.getPublic();
        k = ecpriv.getD();
        c1 = ecpub.getQ();
        SM2keyPair keypair = new SM2keyPair();
        keypair.prvk = k.toString(16);
        keypair.pukX = c1.normalize().getXCoord().toString();
        keypair.pukY = c1.normalize().getYCoord().toString();
        return keypair;
    }

    public ECPoint InitEncipher(ECPoint userKey) {
        BigInteger k = null;
        ECPoint c1 = null;
        AsymmetricCipherKeyPair key = eccKeyPairGenerator.generateKeyPair();
        ECPrivateKeyParameters ecpriv = (ECPrivateKeyParameters) key
                .getPrivate();
        ECPublicKeyParameters ecpub = (ECPublicKeyParameters) key.getPublic();
        k = ecpriv.getD();
        c1 = ecpub.getQ();
        p2 = userKey.multiply(k);
        reset();
        return c1;

    }

    public void Encrypt(byte[] data) {
        sm3c3.update(data, 0, data.length);
        for (int i = 0; i < data.length; i++) {
            if (keyOff == key.length)
                nextKey();

            data[i] ^= key[keyOff++];
        }
    }

    public void InitDecipher(BigInteger userD, ECPoint c1) {
        p2 = c1.multiply(userD);
        reset();
    }

    public void Decrypt(byte[] data) {
        for (int i = 0; i < data.length; i++) {
            if (keyOff == key.length)
                nextKey();

            data[i] ^= key[keyOff++];
        }
        sm3c3.update(data, 0, data.length);
    }

    public void Dofinal(byte[] c3) {
        byte[] p = ByteTool.byteConvert32Bytes(p2.normalize().getYCoord()
                .toBigInteger());
        sm3c3.update(p, 0, p.length);
        sm3c3.doFinal(c3, 0);
        reset();
    }

    /**
     *
     * @param pubKey
     * @param plaintext
     * @param cipherMode
     * @return
     */
    public String Encrypt(ECPoint pubKey, byte[] plaintext, int cipherMode) {
        byte[] data = new byte[plaintext.length];
        System.arraycopy(plaintext, 0, data, 0, plaintext.length);
        ECPoint c1 = InitEncipher(pubKey);
        Encrypt(data);
        byte[] c3 = new byte[32];
        Dofinal(c3);
        String c1Str = StringUtil.padding(c1.normalize().getXCoord()
                .toBigInteger().toString(16), "left", "0", 64)
                + StringUtil.padding(c1.normalize().getYCoord().toBigInteger()
                .toString(16), "left", "0", 64);
        String c2Str = ByteTool.byte2hex(data);
        String c3Str = ByteTool.byte2hex(c3);
        String hexString;
        if (cipherMode == 1) {
            hexString = c1Str + c2Str + c3Str;
        } else {
            hexString = c1Str + c3Str + c2Str;

        }
        // System.out.println("c1-len:" + c1Str.length() + "---" + c1Str);
        // System.out.println("c2-len:" + c2Str.length() + "---" + c2Str);
        // System.out.println("c3-len:" + c3Str.length() + "---" + c3Str);
        return hexString;

    }

    public byte[] Decrypt(BigInteger privateKey, String ciphertext,
                          int cipherMode) {

        String hexString = ciphertext;
        String c1X = hexString.substring(0, 64);
        String c1Y = hexString.substring(0 + c1X.length(), 128);
        String encrypData;
        String c3;
        if (cipherMode == 1) {
            encrypData = hexString.substring(c1X.length() + c1Y.length(),
                    hexString.length() - 64);

            c3 = hexString.substring(hexString.length() - 64,
                    hexString.length());
        } else {
            c3 = hexString.substring(c1X.length() + c1Y.length(), c1X.length()
                    + c1Y.length() + 64);
            encrypData = hexString.substring(
                    c1X.length() + c1Y.length() + c3.length(),
                    hexString.length());
        }
        byte[] data = ByteTool.hexStr2Bytes(encrypData);

        ECPoint c1 = CreatePoint(c1X, c1Y);
        InitDecipher(privateKey, c1);
        Decrypt(data);
        byte[] c3_ = new byte[32];
        Dofinal(c3_);
        // String decryptData = new String(data);
        Boolean isDecrypt = ByteTool.byte2hex(c3_).equalsIgnoreCase(c3);
        return (isDecrypt ? data : new byte[0]);
    }

    public static ECPoint CreatePoint(String x, String y) {

        BigInteger biX = new BigInteger(x, 16);
        BigInteger biY = new BigInteger(y, 16);
        ECPoint point = ecc_curve.createPoint(biX, biY);
        return point;
    }

    public static byte[] sm2GetZ(byte[] userId, ECPoint userKey) {
        SM3Digest sm3 = new SM3Digest();

        int len = userId.length * 8;
        sm3.update((byte) (len >> 8 & 0xFF));
        sm3.update((byte) (len & 0xFF));
        sm3.update(userId, 0, userId.length);

        byte[] p = ByteTool.byteConvert32Bytes(ecc_a);
        sm3.update(p, 0, p.length);

        p = ByteTool.byteConvert32Bytes(ecc_b);
        sm3.update(p, 0, p.length);

        p = ByteTool.byteConvert32Bytes(ecc_gx);
        sm3.update(p, 0, p.length);

        p = ByteTool.byteConvert32Bytes(ecc_gy);
        sm3.update(p, 0, p.length);

        p = ByteTool.byteConvert32Bytes(userKey.normalize().getXCoord()
                .toBigInteger());
        sm3.update(p, 0, p.length);

        p = ByteTool.byteConvert32Bytes(userKey.normalize().getYCoord()
                .toBigInteger());
        sm3.update(p, 0, p.length);

        byte[] md = new byte[sm3.getDigestSize()];
        sm3.doFinal(md, 0);
        return md;
    }

    public static void sm2Sign(byte[] md, BigInteger userD, ECPoint userKey,
                               SM2Result sm2Result) {
        BigInteger e = new BigInteger(1, md);
        BigInteger k = null;
        ECPoint kp = null;
        BigInteger r = null;
        BigInteger s = null;
        do {
            do {
                k = userD;
                kp = userKey;
                System.out.println("计算曲线点X1: "
                        + kp.normalize().getXCoord().toBigInteger()
                        .toString(16));
                System.out.println("计算曲线点Y1: "
                        + kp.normalize().getYCoord().toBigInteger()
                        .toString(16));
                System.out.println("");
                // r
                r = e.add(kp.normalize().getXCoord().toBigInteger());
                r = r.mod(ecc_n);
            } while (r.equals(BigInteger.ZERO) || r.add(k).equals(ecc_n));
            // (1 + dA)~-1
            BigInteger da_1 = userD.add(BigInteger.ONE);
            da_1 = da_1.modInverse(ecc_n);
            // s
            s = r.multiply(userD);
            s = k.subtract(s).mod(ecc_n);
            s = da_1.multiply(s).mod(ecc_n);
        } while (s.equals(BigInteger.ZERO));

        sm2Result.r = r;
        sm2Result.s = s;
    }

    public static void sm2Verify(byte md[], ECPoint userKey, BigInteger r,
                                 BigInteger s, SM2Result sm2Result) {
        sm2Result.R = null;
        BigInteger e = new BigInteger(1, md);
        BigInteger t = r.add(s).mod(ecc_n);
        if (t.equals(BigInteger.ZERO)) {
            return;
        } else {
            ECPoint x1y1 = ecc_point_g.multiply(sm2Result.s);
            System.out.println("计算曲线点X0: "
                    + x1y1.normalize().getXCoord().toBigInteger().toString(16));
            System.out.println("计算曲线点Y0: "
                    + x1y1.normalize().getYCoord().toBigInteger().toString(16));
            System.out.println("");

            x1y1 = x1y1.add(userKey.multiply(t));
            System.out.println("计算曲线点X1: "
                    + x1y1.normalize().getXCoord().toBigInteger().toString(16));
            System.out.println("计算曲线点Y1: "
                    + x1y1.normalize().getYCoord().toBigInteger().toString(16));
            System.out.println("");
            sm2Result.R = e.add(x1y1.normalize().getXCoord().toBigInteger())
                    .mod(ecc_n);
            System.out.println("R: " + sm2Result.R.toString(16));
            return;
        }
    }

    public static byte[] sign(byte[] userId, String privateKey,
                              byte[] sourceData) {
        if (privateKey == null || privateKey.length() == 0) {
            return null;
        }

        if (sourceData == null || sourceData.length == 0) {
            return null;
        }

        BigInteger userD = new BigInteger(privateKey,16);
        System.out.println("userD: " + userD.toString(16));
        System.out.println("");

        ECPoint userKey = ecc_point_g.multiply(userD);
        System.out.println("椭圆曲线点X: "
                + userKey.normalize().getXCoord().toBigInteger().toString(16));
        System.out.println("椭圆曲线点Y: "
                + userKey.normalize().getYCoord().toBigInteger().toString(16));
        System.out.println("");

        SM3Digest sm3 = new SM3Digest();
        byte[] z = sm2GetZ(userId, userKey);
        System.out.println("SM3摘要Z: " + ByteTool.byte2hex(z));
        System.out.println("");

        System.out.println("M: " + ByteTool.byte2hex(sourceData));
        System.out.println("");

        sm3.update(z, 0, z.length);
        sm3.update(sourceData, 0, sourceData.length);
        byte[] md = new byte[32];
        sm3.doFinal(md, 0);

        System.out.println("SM3摘要值: " + ByteTool.byte2hex(md));
        System.out.println("");

        SM2Result sm2Result = new SM2Result();
        sm2Sign(md, userD, userKey, sm2Result);
        System.out.println("r1: " + sm2Result.r.toString(16));
        System.out.println("s1: " + sm2Result.s.toString(16));
        System.out.println("");
        byte[] result = new byte[64];
        byte[] rByte = ByteTool.byteConvert32Bytes(sm2Result.r);
        byte[] sByte = ByteTool.byteConvert32Bytes(sm2Result.s);

        System.arraycopy(rByte, 0, result, 0, 32);
        System.arraycopy(sByte, 0, result, 32, 32);
        return result;
    }


    public static boolean verifySign(byte[] userId, String pubX,String pubY,
                                     byte[] sourceData, byte[] signData) {

        if (sourceData == null || sourceData.length == 0) {
            return false;
        }
        BigInteger biX = new BigInteger(pubX, 16);
        BigInteger biY = new BigInteger(pubY, 16);
        ECPoint userKey = ecc_curve.createPoint(biX, biY);
        SM3Digest sm3 = new SM3Digest();
        byte[] z = sm2GetZ(userId, userKey);
        sm3.update(z, 0, z.length);
        sm3.update(sourceData, 0, sourceData.length);
        byte[] md = new byte[32];
        sm3.doFinal(md, 0);
        System.out.println("SM3摘要值: " + ByteTool.byte2hex(md));
        System.out.println("");
        SM2Result sm2Result = null;
        sm2Result = new SM2Result();
        byte[] rByte = new byte[32];
        byte[] sByte = new byte[32];
        System.arraycopy(signData, 0, rByte, 0, 32);
        System.arraycopy(signData, 32, sByte, 0, 32);
        sm2Result.r = ByteTool.byteConvertInteger(rByte);
        sm2Result.s = ByteTool.byteConvertInteger(sByte);
        System.out.println("r: " + sm2Result.r.toString(16));
        System.out.println("s: " + sm2Result.s.toString(16));
        System.out.println("");
        sm2Verify(md, userKey, sm2Result.r, sm2Result.s, sm2Result);
        return sm2Result.r.equals(sm2Result.R);

    }

}

class SM2Result {
    public BigInteger r;
    public BigInteger s;
    public BigInteger R;
}

